from net import Encoder, Decoder
import sys, os

sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.getcwd()))))
from loss import loss
import torch

batch_size = 1
x = torch.randn(batch_size, 1, 32000)
encoder = Encoder()
decoder = Decoder()
mix_w = encoder(x)
est_mask = torch.randn(2, batch_size, 256, 3999)
mix_w_1 = torch.stack([mix_w] * 2)

y1 = decoder(mix_w_1[0])
y2 = decoder(mix_w_1[1])

print(y1.shape)
print(y2.shape)
device = "cpu"

true = torch.randn(2, 1, 32000).to(device)

y_all = torch.stack([y1, y2], dim=1).permute(1, 0, 2).contiguous().to(device)
print(loss(y_all, true, "cpu"))
print(y_all.shape)
